import argparse
import subprocess
import gc
from typing import List, Tuple, Dict
import cv2
import numpy as np
import onnxruntime
import os
import psutil

# 类别映射
class_mapping = {
    0: "hero",
    1: "monster",
    2: "goods",
    3: "boos",
    4: "door",
    5: "brand",
    6: "again"
}


def check_requirements(package: str) -> None:
    try:
        __import__(package.split('-')[0])
    except ImportError:
        raise ImportError(f"请安装依赖: pip install {package}")


def is_cuda_available():
    try:
        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return result.returncode == 0
    except FileNotFoundError:
        return False


class YOLOv8Detector:
    """YOLOv8 ONNX检测器（无缓存版本）"""

    def __init__(self, onnx_model: str, confidence_thres: float = 0.90, iou_thres: float = 0.5):
        self.onnx_model = onnx_model
        self.confidence_thres = confidence_thres
        self.iou_thres = iou_thres
        self.classes = class_mapping
        self.color_palette = None  # 延迟初始化颜色

        # 模型相关变量
        self.session = None  # ONNX推理会话
        self.input_width = 0  # 模型输入宽度
        self.input_height = 0  # 模型输入高度

        # 图像尺寸变量（单次推理有效，非缓存）
        self.img_height = 0
        self.img_width = 0

    def load_model(self) -> None:
        """加载ONNX模型并初始化推理会话（仅需调用一次）"""
        if self.session is None:
            sess_options = onnxruntime.SessionOptions()
            sess_options.intra_op_num_threads = 4  # 这里可以用你的线程数函数动态设置
            sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
            sess_options.log_severity_level = 3  # 只显示错误
            # 只用CPU推理
            providers = ["CPUExecutionProvider"]
            # providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if is_cuda_available() else [
            #     "CPUExecutionProvider"]
            self.session = onnxruntime.InferenceSession(self.onnx_model, sess_options, providers=providers)
            input_shape = self.session.get_inputs()[0].shape
            self.input_width, self.input_height = input_shape[2], input_shape[3]

    def get_provider_type(self):
        """判断实际使用的推理引擎类型（CPU/GPU）"""
        if self.session is None:
            raise RuntimeError("请先调用load_model()加载模型")

        used_providers = self.session.get_providers()
        # 检查是否使用了CUDAExecutionProvider（GPU加速）
        if any(provider == 'CUDAExecutionProvider' for provider in used_providers):
            return True
        else:
            return False

    def letterbox(self, img: np.ndarray, new_shape: Tuple[int, int] = (640, 640)) -> Tuple[np.ndarray, Tuple[int, int]]:
        """图像缩放+填充（保持比例）"""
        shape = img.shape[:2]
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
        dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2

        if shape[::-1] != new_unpad:
            img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_AREA)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)), (
        top, left)

    def preprocess(self, img: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int]]:
        """预处理图像（缩放、归一化）"""
        self.img_height, self.img_width = img.shape[:2]
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_resized, pad = self.letterbox(img_rgb, (self.input_width, self.input_height))

        img_data = img_resized.astype(np.float32)
        img_data /= 255.0
        img_data = np.transpose(img_data, (2, 0, 1))
        img_data = np.expand_dims(img_data, axis=0)
        return np.ascontiguousarray(img_data), pad

    def postprocess(self, img: np.ndarray, output: List[np.ndarray], pad: Tuple[int, int]) -> Tuple[
        np.ndarray, List[Dict]]:
        """后处理（仅解析输出+NMS，去除绘图，提高速度）"""
        outputs = np.transpose(np.squeeze(output[0]))
        gain = min(self.input_height / self.img_height, self.input_width / self.img_width)

        # 位置调整
        outputs[:, 0] -= pad[1]  # x
        outputs[:, 1] -= pad[0]  # y

        # 提取置信度分数
        class_scores = outputs[:, 4:]
        max_scores = np.max(class_scores, axis=1)
        class_ids = np.argmax(class_scores, axis=1)

        # 过滤低置信度
        mask = max_scores >= self.confidence_thres
        outputs = outputs[mask]
        max_scores = max_scores[mask]
        class_ids = class_ids[mask]

        detections = []
        boxes = []

        for i in range(len(outputs)):
            x, y, w, h = outputs[i][:4]
            left = int((x - w / 2) / gain)
            top = int((y - h / 2) / gain)
            width = int(w / gain)
            height = int(h / gain)
            center_x = left + width / 2
            center_y = top + height / 2
            bottom_y = top + height
            boxes.append([left, top, width, height])

            detections.append({
                "class_id": int(class_ids[i]),
                "class_name": self.classes[int(class_ids[i])],
                "confidence": float(max_scores[i]),
                "box": [left, top, width, height],
                "center": [center_x, center_y],
                "bottom": bottom_y
            })

        # 应用NMS
        if boxes:
            indices = cv2.dnn.NMSBoxes(boxes, max_scores.tolist(), self.confidence_thres, self.iou_thres)
            if len(indices.shape) > 1:
                indices = indices.flatten()
            detections = [detections[i] for i in indices]

        return img, detections

    def draw_detections(self, img: np.ndarray, box: List[float], score: float, class_id: int) -> None:
        """绘制检测框和标签"""
        x1, y1, w, h = box
        color = self.color_palette[class_id]
        cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
        label = f"{self.classes[class_id]}: {score:.2f}"
        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        label_x, label_y = x1, y1 - 10 if y1 > label_height else y1 + label_height
        cv2.rectangle(
            img, (label_x, label_y - label_height), (label_x + label_width, label_y), color, cv2.FILLED
        )
        cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

    def infer(self, img: np.ndarray) -> Tuple[np.ndarray, List[Dict]]:
        """推理OpenCV截图（直接返回结果，无缓存）"""
        if self.session is None:
            raise RuntimeError("请先调用load_model()加载模型")
        if img is None or not isinstance(img, np.ndarray):
            raise ValueError("输入必须是有效的OpenCV图像数组（np.ndarray）")
        img_data, pad = self.preprocess(img)
        outputs = self.session.run(None, {self.session.get_inputs()[0].name: img_data})
        # 后处理直接返回结果（无缓存）
        output_img, detections = self.postprocess(img, outputs, pad)
        del img_data, outputs
        gc.collect()
        return output_img, detections
